/*
 * Copyright (c) 2023 Huawei Technologies Co.,Ltd.
 *
 * openGauss is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *
 *          http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package demo;

import org.apache.log4j.Logger;
import org.junit.Test;
import org.opengauss.jdbc.PgConnection;
import org.opengauss.jdbc.PgStatement;
import org.opengauss.quickautobalance.Cluster;
import org.opengauss.quickautobalance.ConnectionManager;
import org.opengauss.quickautobalance.DataNode;
import org.opengauss.quickautobalance.ReflectUtil;
import org.opengauss.util.HostSpec;
import org.opengauss.util.PSQLException;
import util.ConnectionParams;

import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * 这个demo是旧的
 */
public class QuickAutoBalanceTest {
    
//    private Log LOGGER = Logger.getLogger(QuickAutoBalanceTest.class.getName());
    private Logger LOGGER = Logger.getLogger(QuickAutoBalanceTest.class);
    
    @Test
    public void quickAutoBalanceTest () throws InterruptedException, SQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams();
        System.out.println(url);
        List<PgConnection> connections = new ArrayList<>();
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connections.add(pgConnection);
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
//            System.out.println(socketAddress);
        }
        LOGGER.info("关闭节点前:" + connectionSize);
        // 在此期间关闭数据库
        Thread.sleep(50 * 1000);
        // 等待20秒后重启数据库
        while (true) {
            for (Iterator<PgConnection> iterator  = connections.iterator(); iterator.hasNext();) {
                PgConnection pgConnection = iterator.next();
                if (!pgConnection.isValid(1)) {
                    String socketAddress = pgConnection.getSocketAddress().split("/")[1];
                    connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) - 1);
                    iterator.remove();
                }
            }
            // 前几次查询到的数据，如果此时还没有触发快速负载均衡，应该查到的是100，0，100
//            System.out.println("清理失效节点后:" + connectionSize.toString());
            LOGGER.info("清理失效节点后:" + connectionSize);
            Thread.sleep(5 * 1000);
        }
    }
    
    // 预期所有连接都被杀掉
    @Test
    public void quickAutoBalanceAllIdleTest () throws InterruptedException, PSQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams();
        System.out.println(url);
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
        }
        LOGGER.info("关闭节点前:" + connectionSize);
        // 在此期间关闭数据库
        // 等待20秒后重启数据库
        while (true) {
            Map<String, Cluster> cachedClusters = ReflectUtil.getField(ConnectionManager.class, ConnectionManager.getInstance(), Map.class, "cachedClusters");
            String URLIdentifier = ConnectionManager.getURLIdentifierFromUrl(url);
            Cluster cluster = cachedClusters.get(URLIdentifier);
            Map<HostSpec, DataNode> cachedDnList = ReflectUtil.getField(Cluster.class, cluster, Map.class, "cachedDnList");
            Map<HostSpec, Integer> cachedConnectionSize = new HashMap<>();
            for (Map.Entry<HostSpec, DataNode> entry : cachedDnList.entrySet()) {
                HostSpec hostSpec = entry.getKey();
                DataNode dataNode = entry.getValue();
                cachedConnectionSize.put(hostSpec, dataNode.getCachedConnectionListSize());
            }
            LOGGER.info("关闭某个节点后：" + cachedConnectionSize.toString());
            Thread.sleep(10*1000);
        }
    }
    
    // 预期保留40个连接
    @Test
    public void quickAutoBalanceAllIdleWithParamsTest () throws InterruptedException, PSQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams() + "&minReservedConPerCluster=20&minReservedConPerDatanode=30";
        System.out.println(url);
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        int query = 90;
        List<PgConnection> connections = new ArrayList<>();
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);

            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
            if (i < query) {
                connections.add(pgConnection);
            }
        }
        LOGGER.info("关闭节点前:" + connectionSize.toString());
        // 在此期间关闭数据库
        // 等待20秒后重启数据库
        while (true) {
            Map<String, Cluster> cachedClusters = ReflectUtil.getField(ConnectionManager.class, ConnectionManager.getInstance(), Map.class, "cachedClusters");
            String URLIdentifier = ConnectionManager.getURLIdentifierFromUrl(url);
            Cluster cluster = cachedClusters.get(URLIdentifier);
            Map<HostSpec, DataNode> cachedDnList = ReflectUtil.getField(Cluster.class, cluster, Map.class, "cachedDnList");
            Map<HostSpec, Integer> cachedConnectionSize = new HashMap<>();
            for (Iterator<PgConnection> iterator = connections.iterator(); iterator.hasNext();) {
                PgConnection pgConnection = iterator.next();
                try {
                    if (!pgConnection.isValid(1)) {
                         iterator.remove();
                    }
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }
            for (Map.Entry<HostSpec, DataNode> entry : cachedDnList.entrySet()) {
                HostSpec hostSpec = entry.getKey();
                DataNode dataNode = entry.getValue();
                cachedConnectionSize.put(hostSpec, dataNode.getCachedConnectionListSize());
            }
            LOGGER.info("关闭某个节点后：" + cachedConnectionSize.toString());
            Thread.sleep(10*1000);
        }
    }
    
    @Test
    public void quickAutoBalanceSomeIdleWithTest () throws InterruptedException, PSQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams() + "&minReservedConPerCluster=20&minReservedConPerDatanode=30";
        System.out.println(url);
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        for (int i = 0; i < num; i++) {
            PgConnection pgConnection;
            try {
                pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            String socketAddress = pgConnection.getSocketAddress().split("/")[1];
            connectionSize.put(socketAddress, connectionSize.getOrDefault(socketAddress, 0) + 1);
        }
        LOGGER.info("关闭节点前:" + connectionSize.toString());
        // 在此期间关闭数据库
        // 等待20秒后重启数据库
        while (true) {
            Map<String, Cluster> cachedClusters = ReflectUtil.getField(ConnectionManager.class, ConnectionManager.getInstance(), Map.class, "cachedClusters");
            String URLIdentifier = ConnectionManager.getURLIdentifierFromUrl(url);
            Cluster cluster = cachedClusters.get(URLIdentifier);
            Map<HostSpec, DataNode> cachedDnList = ReflectUtil.getField(Cluster.class, cluster, Map.class, "cachedDnList");
            Map<HostSpec, Integer> cachedConnectionSize = new HashMap<>();
            for (Map.Entry<HostSpec, DataNode> entry : cachedDnList.entrySet()) {
                HostSpec hostSpec = entry.getKey();
                DataNode dataNode = entry.getValue();
                cachedConnectionSize.put(hostSpec, dataNode.getCachedConnectionListSize());
            }
            LOGGER.info("关闭某个节点后：" + cachedConnectionSize.toString());
            Thread.sleep(10*1000);
        }
    }
    
    @Test
    public void isValidTest() throws SQLException {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams();
        System.out.println(url);
        List<PgConnection> connections = new ArrayList<>();
        Map<String, Integer> connectionSize = new HashMap<>();
        int num = 300;
        PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
        pgConnection.isValid(1);
        PgStatement pgStatement = (PgStatement) pgConnection.createStatement();
        pgStatement.execute("select 1");
    }
    
    // 测试参数配置，改成参数错误抛异常后，用例失效了。
    @Test
    public void singleConnectionTest() {
        String url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams() + "&minReservedConPerCluster=20&minReservedConPerDatanode=30";
        System.out.println(url);
        try {
            PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            pgConnection.close();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    
        url = ConnectionParams.getUrlWithCluster() + ConnectionParams.getQuickAutoBalanceParams() + "&minReservedConPerCluster=aaa&minReservedConPerDatanode=aaa";
        System.out.println(url);
        try {
            PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            pgConnection.close();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    
        url = ConnectionParams.getUrlWithCluster() + "autoBalance=leastconn&enableQuickAutoBalance=true&minIdleTimeBeforeTerminal=10&loggerLevel=INFO"
            + "&minReservedConPerCluster=-2&minReservedConPerDatanode=-10";
        System.out.println(url);
        try {
            PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            pgConnection.close();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
        
        url = ConnectionParams.getUrlWithCluster() + "autoBalance=sasasa&enableQuickAutoBalance=sdsds&minIdleTimeBeforeTerminal=10&loggerLevel=INFO"
            + "&minReservedConPerCluster=aaa&minReservedConPerDatanode=aaa";
        System.out.println(url);
        try {
            PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            pgConnection.close();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }
    
    @Test
    public void setParamsTest() {
        String url = ConnectionParams.getUrlWithCluster() + "autoBalance=sasasa&enableQuickAutoBalance=sdsds&minIdleTimeBeforeTerminal=10&loggerLevel=INFO"
            + "&minReservedConPerCluster=aaa&minReservedConPerDatanode=aaa";
        System.out.println(url);
        try {
            PgConnection pgConnection = (PgConnection) DriverManager.getConnection(url, ConnectionParams.user, ConnectionParams.password);
            pgConnection.close();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }
}
